{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ec8bc37",
   "metadata": {},
   "source": [
    "### Test code to count SV carriers "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c54d258",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "import numpy as np\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8da135dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "wkdir = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v3/\"\n",
    "wkdir_path = Path(wkdir)\n",
    "\n",
    "chrom = \"chr21\"\n",
    "sv_info_path = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/lof_missense/data/sv_vcf/info_table/final_sites_critical_info_allele.txt\"\n",
    "vep_msc_path = wkdir_path.joinpath(\"4_vrnt_enrich/sv_vep/msc/SV_vep_hg38_msc_parsed.tsv\")\n",
    "vep_all_path =  wkdir_path.joinpath(\"4_vrnt_enrich/sv_vep/all/SV_vep_hg38_all_parsed.tsv\")\n",
    "z_cutoff_list = [2, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]\n",
    "af_cutoff_list = [0, 0.01, 0.05, 0.1, 0.5]\n",
    "root_dir = wkdir_path.joinpath(\"4_vrnt_enrich/sv_count_carriers/gene_body/200kb_window\")\n",
    "af_cutoff_bins = [(af[1], af_cutoff_list[af[0]+1]) for af in enumerate(af_cutoff_list[:-1])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a484e9b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inputs:\n",
      "- Chromosome: chr21\n",
      "- AF bins: [(0, 0.01), (0.01, 0.05), (0.05, 0.1), (0.1, 0.5)]\n",
      "- Z-score cutoffs: [2, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Inputs:\")\n",
    "print(f\"- Chromosome: {chrom}\")\n",
    "print(f\"- AF bins: {af_cutoff_bins}\")\n",
    "print(f\"- Z-score cutoffs: {z_cutoff_list}\")\n",
    "print(\"\")\n",
    "\n",
    "vep_msc_ranks = [\n",
    "    \"transcript_ablation\",\n",
    "    \"splice_acceptor_variant\",\n",
    "    \"splice_donor_variant\",\n",
    "    \"stop_gained\",\n",
    "    \"frameshift_variant\",\n",
    "    \"stop_lost\",\n",
    "    \"start_lost\",\n",
    "    \"transcript_amplification\",\n",
    "    \"inframe_insertion\",\n",
    "    \"inframe_deletion\",\n",
    "    \"missense_variant\",\n",
    "    \"protein_altering_variant\",\n",
    "    \"splice_donor_5th_base_variant\", \n",
    "    \"splice_region_variant\",\n",
    "    \"splice_donor_region_variant\", \n",
    "    \"splice_polypyrimidine_tract_variant\"\n",
    "    \"incomplete_terminal_codon_variant\",\n",
    "    \"start_retained_variant\",\n",
    "    \"stop_retained_variant\",\n",
    "    \"synonymous_variant\",\n",
    "    \"coding_sequence_variant\",\n",
    "    \"mature_miRNA_variant\",\n",
    "    \"5_prime_UTR_variant\",\n",
    "    \"3_prime_UTR_variant\",\n",
    "    \"non_coding_transcript_exon_variant\",\n",
    "    \"intron_variant\",\n",
    "    \"NMD_transcript_variant\",\n",
    "    \"non_coding_transcript_variant\",\n",
    "    \"coding_transcript_variant\",\n",
    "    \"upstream_gene_variant\",\n",
    "    \"downstream_gene_variant\",\n",
    "    \"TFBS_ablation\",\n",
    "    \"TFBS_amplification\",\n",
    "    \"TF_binding_site_variant\",\n",
    "    \"regulatory_region_ablation\",\n",
    "    \"regulatory_region_amplification\",\n",
    "    \"feature_elongation\",\n",
    "    \"regulatory_region_variant\",\n",
    "    \"feature_truncation\",\n",
    "    \"intergenic_variant\",\n",
    "    \"sequence_variant\",\n",
    "    \"no_predicted_effect\"]\n",
    "\n",
    "# from all VEP consequences that have no gene annotation\n",
    "vep_msc_not_linked_to_gene = ['TFBS_ablation', \n",
    "                              'TF_binding_site_variant',\n",
    "                              'regulatory_region_variant', \n",
    "                              'TFBS_amplification',\n",
    "                              'intergenic_variant', \n",
    "                              'regulatory_region_ablation',\n",
    "                              'regulatory_region_amplification']\n",
    "\n",
    "root_dir_path = Path(root_dir)\n",
    "tpm_cutoff = 0.5\n",
    "\n",
    "# load gene expression matrix \n",
    "ge_matrix_flat_chrom_egan_path = root_dir_path.joinpath(f\"express_mtx/{chrom}_ge_mtx_flat.tsv\")\n",
    "ge_matrix_flat_chrom_egan_df = pd.read_csv(ge_matrix_flat_chrom_egan_path, sep=\"\\t\")\n",
    "# add gene-sample pairs\n",
    "ge_matrix_flat_chrom_egan_df[\"gene_smpl_pair\"] = ge_matrix_flat_chrom_egan_df.gene_id + \",\" + ge_matrix_flat_chrom_egan_df.rna_id\n",
    "# carrier information \n",
    "sv_intersect_express_info_path = root_dir_path.joinpath(f\"express_carrier_info/{chrom}_express_carrier_info.tsv\")\n",
    "sv_intersect_express_info_df = pd.read_csv(sv_intersect_express_info_path, sep=\"\\t\")\n",
    "# add gene-sample pairs  \n",
    "sv_intersect_express_info_df[\"gene_smpls_id\"] = sv_intersect_express_info_df.gene_id.astype(str) + \",\" + sv_intersect_express_info_df.rna_id.astype(str)\n",
    "# RNA IDs with EGAN ID match \n",
    "sv_intersect_express_info_path = root_dir_path.joinpath(\"egan_rna_smpls/egan_rna_ids_paired_pass_qc.tsv\")\n",
    "rna_id_pass_qc_sv_calls = pd.read_csv(sv_intersect_express_info_path, sep=\"\\t\").rna_id.unique()\n",
    "# SV types \n",
    "sv_info_df = pd.read_csv(sv_info_path, sep=\"\\t\", dtype={\"plinkID\":str})\n",
    "sv_info_id_af_df = sv_info_df[[\"plinkID\", \"AF\", \"SVTYPE\"]].rename(columns={\"plinkID\":\"vrnt_id\"})\n",
    "sv_types_list = sv_info_df.SVTYPE.unique().tolist()\n",
    "\n",
    "# annotate variant-gene pairs with VEP consequence \n",
    "vep_all_consq_df = pd.read_csv(vep_all_path, sep=\"\\t\")\n",
    "# variant gene pairs \n",
    "vrnt_gene_pairs_df = sv_intersect_express_info_df[[\"vrnt_id\", \"gene_id\"]].drop_duplicates()\n",
    "# collapse to gene-variant consequence pairs \n",
    "vep_all_consq_drop_dups_df = vep_all_consq_df[[\"vrnt_id\", \"gene_id\", \"Consequence\"]].drop_duplicates()\n",
    "# add variant-gene consequences \n",
    "vrnt_gene_pairs_consq_df = pd.merge(vrnt_gene_pairs_df, \n",
    "                                    vep_all_consq_drop_dups_df, \n",
    "                                    on=[\"vrnt_id\", \"gene_id\"], \n",
    "                                    how=\"left\").fillna(\"no_predicted_effect\")\n",
    "# group consequences together for each variant-gene pair\n",
    "vrnt_gene_pairs_consq_df[\"gene_consequence\"] = vrnt_gene_pairs_consq_df.groupby([\"vrnt_id\", \"gene_id\"])['Consequence'].transform(lambda x: ','.join(x))\n",
    "# remove duplicates arising from variant have different consequences across gene transcripts \n",
    "vrnt_gene_pair_consq_collapse_df = vrnt_gene_pairs_consq_df.drop(columns=[\"Consequence\"]).drop_duplicates()\n",
    "if vrnt_gene_pair_consq_collapse_df.shape[0] != vrnt_gene_pairs_df.shape[0]: \n",
    "    raise ValueError(\"Variant gene pair number does not match number of variant gene pairs with consequence.\")\n",
    "# assign each variant-gene pair a unique consequence based on VEP rank  \n",
    "vrnt_gene_pair_msc_consq = []\n",
    "for index, row in vrnt_gene_pair_consq_collapse_df.iterrows(): \n",
    "    gene_consequence = row[\"gene_consequence\"]\n",
    "    gene_consequence_list = gene_consequence.split(\",\")\n",
    "    for consq in vep_msc_ranks: \n",
    "        if consq in gene_consequence_list: \n",
    "            break \n",
    "    vrnt_gene_pair_msc_consq.append(consq)\n",
    "if vrnt_gene_pair_consq_collapse_df.shape[0] != len(vrnt_gene_pair_msc_consq): \n",
    "    raise ValueError(\"Number of MSC per gene does not match number of variant-gene pairs.\")\n",
    "vrnt_gene_pair_consq_collapse_df[\"consequence\"] = vrnt_gene_pair_msc_consq\n",
    "vrnt_gene_pair_consq_collapse_df = vrnt_gene_pair_consq_collapse_df.drop(columns=[\"gene_consequence\"])\n",
    "# split into variants with annotated gene effect and no predicted gene effect\n",
    "vrnt_gene_effect_df = vrnt_gene_pair_consq_collapse_df[vrnt_gene_pair_consq_collapse_df.consequence != \"no_predicted_effect\"]\n",
    "no_predicted_effect_df = vrnt_gene_pair_consq_collapse_df[vrnt_gene_pair_consq_collapse_df.consequence == \"no_predicted_effect\"]\n",
    "# load VEP MSC annotations\n",
    "vep_msc_df = pd.read_csv(vep_msc_path, sep=\"\\t\").rename(columns={\"Uploaded_variation\": \"vrnt_id\", \"Consequence\": \"msc\"})\n",
    "# merge VEP MSC with no predicted effect \n",
    "no_prediced_effect_df = pd.merge(no_predicted_effect_df[[\"vrnt_id\", \"gene_id\"]], \n",
    "                                 vep_msc_df[[\"vrnt_id\", \"msc\"]], \n",
    "                                 on=\"vrnt_id\", \n",
    "                                 how=\"left\")\n",
    "# check for NaNs \n",
    "if no_prediced_effect_df.msc.isnull().values.any(): \n",
    "    raise ValueError(\"Variants missing VEP most severe consequence.\")\n",
    "no_prediced_effect_df[\"consequence\"] = np.where(no_prediced_effect_df.msc.isin(vep_msc_not_linked_to_gene), \n",
    "                                                no_prediced_effect_df.msc, \n",
    "                                                \"no_predicted_effect\")\n",
    "no_prediced_effect_trunc_df = no_prediced_effect_df[[\"vrnt_id\", \"gene_id\", \"consequence\"]]\n",
    "# combine variants with annotated gene effect and variants with no predicted effect with updated regulatory consequence\n",
    "vrnt_gene_pair_consq_msc_added_df = pd.concat([vrnt_gene_effect_df, no_prediced_effect_trunc_df])\n",
    "if vrnt_gene_pair_consq_msc_added_df[[\"vrnt_id\", \"gene_id\"]].drop_duplicates().shape[0] != vrnt_gene_pairs_df.shape[0]: \n",
    "    raise ValueError(\"Number of gene pairs in expression input does not match number of gene pairs with annotated consequence.\")\n",
    "# add consequence to carrier information \n",
    "sv_intersect_express_info_gene_msc_df = pd.merge(sv_intersect_express_info_df, \n",
    "                                                 vrnt_gene_pair_consq_msc_added_df, \n",
    "                                                 on=[\"vrnt_id\", \"gene_id\"], \n",
    "                                                 how=\"left\")\n",
    "# check number of rows not changed \n",
    "if sv_intersect_express_info_gene_msc_df.shape[0] != sv_intersect_express_info_df.shape[0]: \n",
    "    raise ValueError(\"Added entries to sample-gene-variant expression dataframe.\")\n",
    "# check for nans \n",
    "if sv_intersect_express_info_gene_msc_df.consequence.isnull().values.any(): \n",
    "    raise ValueError(\"Missing consequence annotations for some variant-gene pairs.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "439982a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "### count carriers in test and controls \n",
    "carrier_count = {}\n",
    "gene_smpl_pair_carrier = {}\n",
    "gene_smpl_count = 0 \n",
    "for af_range in af_cutoff_bins:\n",
    "    af_lower, af_upper = af_range\n",
    "    for z_cutoff in z_cutoff_list: \n",
    "        carrier_count[gene_smpl_count] = [chrom, z_cutoff, f\"{round(af_lower*100)}-{round(af_upper*100)}\"]\n",
    "        ### get control and test gene-sample pairs based on misexpression cutoff \n",
    "        # get gene-sample pairs passing z-score cutoff and misexpression cutoff\n",
    "        z_cutoff_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) & \n",
    "                                                   (ge_matrix_flat_chrom_egan_df[\"TPM\"] > tpm_cutoff)\n",
    "                                                  ]\n",
    "        test_gene_ids = z_cutoff_df.gene_id.unique()\n",
    "        test_gene_smpl_pairs = z_cutoff_df.gene_smpl_pair.unique()\n",
    "        # get gene-sample pairs in control group \n",
    "        cntrl_gene_smpl_pairs = ge_matrix_flat_chrom_egan_df[~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(test_gene_smpl_pairs) &\n",
    "                                                             ge_matrix_flat_chrom_egan_df.gene_id.isin(test_gene_ids)\n",
    "                                                            ].gene_smpl_pair.unique()\n",
    "        # check overlap between sets is empty \n",
    "        if len(set(test_gene_smpl_pairs).intersection(set(cntrl_gene_smpl_pairs))) != 0:\n",
    "            raise ValueError(\"Overlap between control and test gene-sample pair sets\")\n",
    "        # check number of test and control gene-pairs matches expected \n",
    "        total_test_control_pairs = len(test_gene_smpl_pairs) + len(cntrl_gene_smpl_pairs)\n",
    "        test_genes_by_smpls = len(test_gene_ids) * len(rna_id_pass_qc_sv_calls)\n",
    "        if total_test_control_pairs !=  test_genes_by_smpls: \n",
    "            raise ValueError(\"Number of test and control gene pairs does not match test gene IDs by samples\")\n",
    "        carrier_count[gene_smpl_count] += [len(test_gene_ids), len(rna_id_pass_qc_sv_calls), len(test_gene_smpl_pairs), len(cntrl_gene_smpl_pairs)]\n",
    "\n",
    "        ### count carriers in control and test groups \n",
    "        # get misexpression carriers \n",
    "        misexpress_carriers_df = sv_intersect_express_info_gene_msc_df[(sv_intersect_express_info_gene_msc_df.gene_smpls_id.isin(test_gene_smpl_pairs)) & \n",
    "                                                              (sv_intersect_express_info_gene_msc_df.AF >= af_lower) & \n",
    "                                                              (sv_intersect_express_info_gene_msc_df.AF < af_upper) &\n",
    "                                                              (sv_intersect_express_info_gene_msc_df.genotype.isin(['(0, 1)', '(1, 1)']))\n",
    "                                                             ].copy()\n",
    "        # get non-misexpression carriers\n",
    "        cntrl_carriers_df = sv_intersect_express_info_gene_msc_df[(sv_intersect_express_info_gene_msc_df.gene_smpls_id.isin(cntrl_gene_smpl_pairs)) &\n",
    "                                                                  (sv_intersect_express_info_gene_msc_df.AF >= af_lower) & \n",
    "                                                                  (sv_intersect_express_info_gene_msc_df.AF < af_upper) &\n",
    "                                                                  (sv_intersect_express_info_gene_msc_df.genotype.isin(['(0, 1)', '(1, 1)']))\n",
    "                                                                 ].copy()\n",
    "        # add number of carriers \n",
    "        gene_smpl_misexpress_carriers = misexpress_carriers_df.gene_smpls_id.unique()\n",
    "        gene_smpl_cntrl_carriers = cntrl_carriers_df.gene_smpls_id.unique()\n",
    "        carrier_count[gene_smpl_count] += [len(gene_smpl_misexpress_carriers), len(gene_smpl_cntrl_carriers)]\n",
    "        # check overlap between sets is empty \n",
    "        if len(set(gene_smpl_misexpress_carriers).intersection(set(gene_smpl_cntrl_carriers))) != 0:\n",
    "            raise ValueError(\"Overlap between control and test gene-sample pair carrier sets\")\n",
    "\n",
    "        ### count carriers for SV types\n",
    "        for sv_type in sv_types_list:\n",
    "            misexpress_carriers_sv_type_df = misexpress_carriers_df[misexpress_carriers_df.SVTYPE == sv_type]\n",
    "            cntrl_carriers_sv_type_df = cntrl_carriers_df[cntrl_carriers_df.SVTYPE == sv_type]\n",
    "            sv_type_misexp_carriers = len(misexpress_carriers_sv_type_df.gene_smpls_id.unique())\n",
    "            sv_type_cntrl_carriers = len(cntrl_carriers_sv_type_df.gene_smpls_id.unique())\n",
    "            carrier_count[gene_smpl_count] += [sv_type_misexp_carriers, sv_type_cntrl_carriers]\n",
    "            for msc in vep_msc_ranks: \n",
    "                msc_type_misexp_carriers = len(misexpress_carriers_sv_type_df[misexpress_carriers_sv_type_df.consequence == msc].gene_smpls_id.unique())\n",
    "                msc_type_cntrl_carriers = len(cntrl_carriers_sv_type_df[cntrl_carriers_sv_type_df.consequence == msc].gene_smpls_id.unique())\n",
    "                carrier_count[gene_smpl_count] += [msc_type_misexp_carriers, msc_type_cntrl_carriers]\n",
    "        gene_smpl_count += 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "19822f02",
   "metadata": {},
   "outputs": [],
   "source": [
    "vrnt_type_cols = []\n",
    "for sv_type in sv_types_list: \n",
    "    vrnt_type_cols += [f\"{sv_type}_misexp\", f\"{sv_type}_contrl\"]\n",
    "    for msc in vep_msc_ranks:\n",
    "        vrnt_type_cols += [f\"{sv_type}_{msc}_misexp\", f\"{sv_type}_{msc}_contrl\"]\n",
    "\n",
    "columns = [\"chrom\", \"z_cutoff\", \"maf_bin\", \"misexp_genes\", \"smpls_pass_qc\", \"total_misexp\", \"total_control\", \"all_sv_misexp\", \"all_sv_contrl\"] + vrnt_type_cols\n",
    "carrier_count_df = pd.DataFrame.from_dict(carrier_count, orient=\"index\", columns=columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9d092b9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
